"""
python lstm_readJson_test.py
"""

import pandas as pd
import numpy as np
import os
import json
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import LSTM, Dense
from tensorflow.keras.preprocessing.sequence import TimeseriesGenerator


def predict_single_category(input_path, output_path):
    # 读取输入的 JSON 文件
    with open(input_path, 'r') as f:
        data = json.load(f)

    # 将 JSON 数据转换为 DataFrame
    df = pd.DataFrame(data)
    df['ds'] = pd.to_datetime(df['date'])
    df.rename(columns={'sales': 'y'}, inplace=True)

    # 准备数据用于训练
    values = df['y'].values
    sequence_length = 10
    generator = TimeseriesGenerator(values, values, length=sequence_length, batch_size=1)

    # 创建 LSTM 模型
    model = Sequential()
    model.add(LSTM(50, input_shape=(sequence_length, 1)))
    model.add(Dense(1))

    # 编译模型
    model.compile(loss='mse', optimizer='adam')

    # 拟合模型
    model.fit(generator, epochs=50, verbose=0)

    # 进行预测
    last_sequence = values[-sequence_length:]
    predictions = []
    for _ in range((pd.to_datetime('2023-09-30') - pd.to_datetime(df['ds'].iloc[-1])).days):
        new_prediction = model.predict(last_sequence.reshape(1, sequence_length, 1))[0][0]
        predictions.append(new_prediction)
        last_sequence = np.append(last_sequence[1:], new_prediction)

    # 创建包含预测日期的数据框
    future_dates = pd.date_range(start=df['ds'].iloc[-1] + pd.Timedelta(days=1), end='2023-09-30')
    result_df = pd.DataFrame({'ds': future_dates, 'yhat': predictions})

    # 将结果转换回 JSON 格式
    result_json = result_df[['ds', 'yhat']].to_json(orient='records')
    result_json = json.loads(result_json)

    # 创建输出目录如果不存在
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 写入输出 JSON 文件
    with open(output_path, 'w') as f:
        json.dump(result_json, f)

    # 绘制原始数据和预测数据的散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(df['ds'], df['y'], label='Original Data', c='blue')
    plt.scatter(result_df['ds'], result_df['yhat'], label='Predicted Data', c='red')
    plt.legend()
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title('LSTM Prediction on Sales Data')

    # 保存图像
    plt.savefig(os.path.join(output_dir, os.path.basename(input_path).replace('.json', '.png')))


if __name__ == "__main__":
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "./unit_test/category_category1_output.json"
    predict_single_category(input_path, output_path)
